import numpy as np
from astropy.time import Time, TimeDelta
from astropy.coordinates import Longitude, Latitude, Angle
from astropy import coordinates
import astropy.units as u
import utils

a_e = 6378137
GE = 3986004418e5
time_unit = np.sqrt(a_e ** 3 / GE)


class InitialOrbit:

    def __init__(self, which_day, datas):
        self.datas = datas
        self.which_day = which_day
        self.F = 0
        self.G = 0
        self.site_lon = 120 * u.deg
        self.site_lat = 36 * u.deg
        self.site_height = 40 * u.m
        self.site_coord_ITRS = coordinates.EarthLocation.from_geodetic(
            self.site_lon, self.site_lat, self.site_height).get_itrs().cartesian.xyz
        self.lambda_mu_nu_tau_R = list()
        self.R = np.zeros(3)
        self.r = np.zeros(3)
        self.r_dot = np.zeros(3)
        self.F = 0
        self.G = 0

    def get_lambda_mu_nu_tau_R(self):
        lambda_mu_nu_tau_R = list()
        for dat in self.datas:
            sec_of_day, azi, alt = dat
            azi = azi.to_value('rad')
            alt = alt.to_value('rad')
            this_dat_time = self.which_day + sec_of_day
            transform_matrix = utils.CoordinatesTransformEquinoxBase(this_dat_time)
            matrix_GR = transform_matrix.nutation_matrix() @ transform_matrix.precession_matrix()
            sg = this_dat_time.sidereal_time('apparent', longitude=self.site_lon).to_value('rad')
            matrix_ZR = utils.rotation_matrix('z', np.pi - sg) @ utils.rotation_matrix('y', np.pi / 2 - np.pi / 5)
            vector_L = matrix_GR.T @ matrix_ZR @ np.array([np.cos(alt) * np.cos(azi), -np.cos(alt) * np.sin(azi), np.sin(alt)])
            this_dat_time_site_GCRS = matrix_GR.T @ transform_matrix.earth_rotation_matrix().T @ self.site_coord_ITRS
            this_dat_time_R = this_dat_time_site_GCRS.to_value('m') / a_e
            lambda_mu_nu_tau_R.append((vector_L, sec_of_day.to_value('s') / time_unit, this_dat_time_R))
            if dat == self.datas[0]:
                self.r = this_dat_time_R
        self.lambda_mu_nu_tau_R = lambda_mu_nu_tau_R

    def update_solution(self):
        A = np.zeros((3 * len(self.lambda_mu_nu_tau_R), 6))
        B = np.zeros(3 * len(self.lambda_mu_nu_tau_R))
        t_0 = 0

        vector_r = self.r
        vector_r_dot = self.r_dot
        scalar_r = np.linalg.norm(vector_r)
        scalar_r_dot = np.linalg.norm(vector_r_dot)
        for i in range(len(self.lambda_mu_nu_tau_R)):
            # 第（4）问，把下面两行注释去掉就可以看到结果
            # if i == 4:
            #     break
            vector_L, sec_of_day, R = self.lambda_mu_nu_tau_R[i]
            if i == 0:
                t_0 = sec_of_day
                # continue
            lmda, mu, nu = vector_L
            X, Y, Z = R
            tau = (sec_of_day - t_0) ** np.arange(8)
            u = (1 / scalar_r) ** np.arange(10)
            sigma = np.dot(vector_r, vector_r_dot)
            v_0_2 = scalar_r_dot ** 2
            # 本来是抄书上的F公式，结果眼都看花了，把z和z_dot都写上去了，结果就是算出来卫星的轨道在地表以下，a高达3216553.0213314m
            # 现在写的是PPT上的公式
            F = (1 - tau[2] / 2 * u[3] + tau[3] / 2 * u[5] * sigma +
                 tau[4] / 24 * u[5] * (3 * v_0_2 - 2 * u[1] - 15 * u[2] * (sigma ** 2)) +
                 tau[5] / 8 * u[7] * 15 * sigma * (-3 * v_0_2 + 2 * u[1] + 7 * u[2] * (sigma ** 2)) +
                 tau[6] / 720 * u[7] * (u[2] * (sigma ** 2) * (630 * v_0_2 - 420 * u[1] - 945 * u[2] * (sigma ** 2)) - (22 * u[2] - 66 * u[1] * v_0_2 + 45 * (v_0_2 ** 2))))

            G = (tau[1] - tau[3] / 6 * u[3] + tau[4] / 4 * u[5] * sigma +
                 tau[5] / 120 * u[5] * (9 * v_0_2 - 8 * u[1] - 45 * u[2] * (sigma ** 2)) +
                 tau[6] / 24 * u[7]  * sigma * (-6 * v_0_2 + 5 * u[1] + 14 * u[2] * (sigma ** 2)))

            A[3 * i + 0] = np.array([F * nu, 0, -F * lmda, G * nu, 0, -G * lmda])
            A[3 * i + 1] = np.array([0, F * nu, -F * mu, 0, G * nu, -G * mu])
            A[3 * i + 2] = np.array([F * mu, -F * lmda, 0, G * mu, -G * lmda, 0])
            B[3 * i + 0] = nu * X - lmda * Z
            B[3 * i + 1] = nu * Y - mu * Z
            B[3 * i + 2] = mu * X - lmda * Y

        solution = np.linalg.lstsq(A, B, rcond=None)
        self.F = F
        self.G = G
        self.r = solution[0][:3]
        self.r_dot = solution[0][3:]


# 第（3）问直接改这里的时间就可以
time_str = '2002-03-30T00:00:00'
utc_time = Time(time_str, scale='utc')
data_list = utils.read_txt('data.txt')

initial_orbit = InitialOrbit(utc_time, data_list)
initial_orbit.get_lambda_mu_nu_tau_R()

# 设置了100次循环，发现差不多第10几次左右这个F和G就变化的特别小了。
for i in range(100):
    initial_orbit.update_solution()
    # initial_orbit.update_F_G()

orbit_elements = utils.get_orbit_elements(initial_orbit.r * a_e, initial_orbit.r_dot * a_e / time_unit)
print(orbit_elements)

# t_0_utc = utc_time + 28470.9946 * u.s
# transform_matrix = utils.CoordinatesTransformEquinoxBase(t_0_utc)
# tete_r = transform_matrix.nutation_matrix() @ transform_matrix.precession_matrix() @ initial_orbit.r
# tete_r_dot = transform_matrix.nutation_matrix() @ transform_matrix.precession_matrix() @ initial_orbit.r_dot
# orbit_elements_tete = utils.get_orbit_elements(tete_r * a_e, tete_r_dot * a_e / time_unit)
# print(orbit_elements_tete)
